package com.gitee.wsl.data.pool

import com.gitee.wsl.ext.coroutines.get
import kotlinx.atomicfu.atomic
import kotlinx.coroutines.TimeoutCancellationException
import kotlinx.coroutines.channels.Channel
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.withContext
import kotlinx.coroutines.withTimeout
import kotlin.coroutines.CoroutineContext
import kotlin.coroutines.coroutineContext
import kotlin.time.Duration.Companion.seconds

interface ConnectionPool<R: AutoCloseable>:AutoCloseable {

    suspend fun useConnection(isReadOnly: Boolean, block: suspend (R) -> R): R

    /**
     * Closes the pool and any opened connections, attempting to use connections is an error once
     * the pool is closed.
     */
    //fun close()
    companion object{
        fun<C: AutoCloseable> create(connectionFactory: () -> C): ConnectionPool<C> = ConnectionPoolImpl(connectionFactory)

        fun<C: AutoCloseable> create(maxNumOfReaders: Int,
                                     maxNumOfWriters: Int,
                                     connectionFactory: () -> C): ConnectionPool<C> = ConnectionPoolImpl(maxNumOfReaders, maxNumOfWriters, connectionFactory)

        fun<C: AutoCloseable> createAsync(capacity: Int,  connectionFactory: () -> C) = SynchronizedPool(capacity, connectionFactory)
    }
}


internal class ConnectionPoolImpl<C: AutoCloseable> : ConnectionPool<C> {

    private val readers: SynchronizedPool<C>
    private val writers: SynchronizedPool<C>

    //private val threadLocal = ThreadLocal<PooledConnectionImpl<C>>()
    private val threadLocal = atomic<PooledConnectionImpl<C>?>(null)

    private val _isClosed = atomic(false)
    private val isClosed by _isClosed

    // Amount of time to wait to acquire a connection before throwing, Android uses 30 seconds in
    // its pool, so we do too here, but IDK if that is a good number. This timeout is unrelated to
    // the busy handler.
    // TODO: Allow configuration
    private val timeout = 30.seconds

    constructor(connectionFactory: () -> C) {
        this.readers = SynchronizedPool(capacity = 1, connectionFactory = connectionFactory)
        this.writers = readers
    }

    constructor(
        maxNumOfReaders: Int,
        maxNumOfWriters: Int,
        connectionFactory: () -> C
    ) {
        require(maxNumOfReaders > 0) { "Maximum number of readers must be greater than 0" }
        require(maxNumOfWriters > 0) { "Maximum number of writers must be greater than 0" }

        this.readers = SynchronizedPool(capacity = maxNumOfReaders, connectionFactory = connectionFactory)
        this.writers = SynchronizedPool(capacity = maxNumOfWriters, connectionFactory = connectionFactory)
    }

    override suspend fun  useConnection(
        isReadOnly: Boolean,
        block: suspend (C) -> C
    ): C {
        if (isClosed) {
            throw Exception("AutoCloseable pool is closed")
        }
        val confinedConnection = threadLocal.get() ?: coroutineContext[ConnectionElement]?.connectionWrapper as PooledConnectionImpl<C>?
        if (confinedConnection != null) {
            if (!isReadOnly && confinedConnection.isReadOnly) {
                throw Exception("Cannot upgrade connection from reader to writer")
            }
            return if (coroutineContext[ConnectionElement] == null) {
                // Reinstall the connection context element if it is missing. We are likely in
                // a new coroutine but were able to transfer the connection via the thread local.
                withContext(createConnectionContext(confinedConnection)) {
                    block.invoke(confinedConnection.rawConnection)
                }
            } else {
                block.invoke(confinedConnection.rawConnection)
            }
        }
        val pool =
            if (isReadOnly) {
                readers
            } else {
                writers
            }
        val result: C
        var exception: Throwable? = null
        var connection: PooledConnectionImpl<C>? = null
        try {
            val (acquiredConnection, acquireError) = pool.acquireWithTimeout()
            // Always try to create a wrapper even if an error occurs, so it can be recycled.
            connection =
                acquiredConnection?.let {
                    PooledConnectionImpl(
                        delegate = it,
                        isReadOnly = readers !== writers && isReadOnly
                    )
                }
            if (acquireError is TimeoutCancellationException) {
                // TODO: Expose more information on which threads are holding into the connections.
                throw Exception( "Timed out attempting to acquire a connection")
            } else if (acquireError != null) {
                throw acquireError
            }
            requireNotNull(connection)
            result = withContext(createConnectionContext(connection)) { block.invoke(connection.rawConnection) }
        } catch (ex: Throwable) {
            exception = ex
            throw ex
        } finally {
            try {
                connection?.let { usedConnection ->
                    usedConnection.markRecycled()
                    pool.recycle(usedConnection.delegate)
                }
            } catch (error: Throwable) {
                exception?.addSuppressed(error)
            }
        }
        return result
    }

    private suspend inline fun SynchronizedPool<C>.acquireWithTimeout(): Pair<ConnectionWithLock<C>?, Throwable?> {
        // Following async timeout with resources recommendation:
        // https://kotlinlang.org/docs/cancellation-and-timeouts.html#asynchronous-timeout-and-resources
        var connection: ConnectionWithLock<C>? = null
        var exceptionThrown: Throwable? = null
        try {
            withTimeout(timeout) { connection = this@acquireWithTimeout.acquire() }
        } catch (ex: Throwable) {
            exceptionThrown = ex
        }
        return connection to exceptionThrown
    }

    private fun createConnectionContext(connection: PooledConnectionImpl<C>) =
        ConnectionElement(connection) /*+ threadLocal.asContextElement(connection)*/

    // TODO: (b/319657104): Make suspending so pool closes when all connections are recycled.
    override fun close() {
        if (_isClosed.compareAndSet(expect = false, update = true)) {
            readers.close()
            writers.close()
        }
    }
}


class ConnectionWithLock<C: AutoCloseable>(
    val delegate: C,
    private val lock: Mutex = Mutex()
) : AutoCloseable by delegate, Mutex by lock


internal class PooledConnectionImpl<C: AutoCloseable>(
    val delegate: ConnectionWithLock<C>,
    val isReadOnly: Boolean,
    val recycleBlock: ConnectionWithLock<C>.()->Unit={}
)  {

    private val _isRecycled = atomic(false)
    private val isRecycled by _isRecycled

    val rawConnection: C
        get() = delegate.delegate

    fun markRecycled() {
        if (_isRecycled.compareAndSet(expect = false, update = true)) {
            // Perform a rollback in case there is an active transaction so that the connection
            // is in a clean state when it is recycled. We don't know for sure if there is an
            // unfinished transaction, hence we always try the rollback.
            // TODO(b/319627988): Try to *really* check if there is an active transaction with the
            //     C APIs sqlite3_txn_state or sqlite3_get_autocommit and possibly throw an error
            //     if there is an unfinished transaction.
            try {
                //delegate.execSQL("ROLLBACK TRANSACTION")
                recycleBlock(delegate)
            } catch (_: Exception) {
                // ignored
            }
        }
    }
}

private class ConnectionElement(
    val connectionWrapper: PooledConnectionImpl<*>
) : CoroutineContext.Element {

    companion object Key : CoroutineContext.Key<ConnectionElement>

    override val key: CoroutineContext.Key<ConnectionElement>
        get() = ConnectionElement
}

/**
 * Creates a [CoroutineContext.Element] from this thread local that will transfer the data to
 * threads resumed by the coroutine context.
 */
//internal fun <T> ThreadLocal<T>.asContextElement(value: T): CoroutineContext.Element =
//    this.asContextElement(value)


open class SynchronizedPool<C: AutoCloseable>(val capacity: Int, val connectionFactory: () -> C):ConnectionPool<C> {
    private val size = atomic(0)
    private val connections = arrayOfNulls<ConnectionWithLock<C>>(capacity)


    private val channel =
        Channel<ConnectionWithLock<C>>(
            capacity = capacity,
            // Only trySend() is used, but due to high paranoia add an undelivered callback
            onUndeliveredElement = { unusedConnection -> unusedConnection.close() }
        )

    override suspend fun useConnection(isReadOnly: Boolean, block: suspend (C) -> C): C{
        val connect =  acquire()
        return try{
              block(connect.delegate)
        } finally {
            recycle(connect)
        }
    }

    suspend fun acquire(): ConnectionWithLock<C> {
        val receiveResult = channel.tryReceive()
        return if (receiveResult.isSuccess) {
            receiveResult.getOrThrow()
        } else {
            tryOpenNewConnection()
            channel.receive()
        }
    }

    private fun tryOpenNewConnection() {
        val currentSize = size.value
        if (currentSize >= capacity) {
            // Capacity reached
            return
        }
        if (size.compareAndSet(currentSize, currentSize + 1)) {
            val newConnection = ConnectionWithLock(connectionFactory.invoke())
            val sendResult = channel.trySend(newConnection)
            if (sendResult.isSuccess) {
                connections[currentSize] = newConnection
            } else {
                newConnection.close()
                if (!sendResult.isClosed) {
                    // Failed to send but channel is not closed, this means a race condition with
                    // the size and capacity checks.
                    error("Couldn't send a new connection for acquisition")
                }
            }
        } else {
            // Another thread went ahead and created a new connection, try again
            tryOpenNewConnection()
        }
    }

    fun recycle(connection: ConnectionWithLock<C>) {
        val sendResult = channel.trySend(connection)
        if (!sendResult.isSuccess) {
            connection.close()
            if (!sendResult.isClosed) {
                // Failed to send but channel is not closed. Likely a race condition...
                // did open connections exceeded capacity? Maybe a `finally` block didn't run?
                error("Couldn't recycle connection")
            }
        }
    }

    override fun close() {
        channel.close()
        connections.forEach { it?.close() }
    }
}